package org.basis.algorithm.tree.common;

/**
 * 前缀树
 *
 * @author Mr_wenpan@163.com 2021/12/24 14:32
 */
public class Trie {
    /**
     * 根节点
     */
    TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public Trie(TrieNode root) {
        this.root = root;
    }

    public String search2(String word, int n) {
        if (word == null) {
            return "";
        }
        char[] chars = word.toCharArray();
        TrieNode cur = root;
        int res = 0;
        for (char aChar : chars) {
            int index = aChar - 'a';
            if (cur.nexts[index] != null && cur.nexts[index].pass == n) {
                res++;
            } else {
                break;
            }
            cur = cur.nexts[index];
        }
        return word.substring(0, res);
    }

    /**
     * 往前缀树插入一个字符串
     */
    public void insert(String word) {
        if (word == null) {
            return;
        }
        TrieNode cur = root;
        char[] chars = word.toCharArray();
        for (char aChar : chars) {
            // 获取到下标
            int index = aChar - 'a';
            // 没有走向下一个字符的路，则新建出来
            if (cur.nexts[index] == null) {
                cur.nexts[index] = new TrieNode();
            }
            cur.pass++;
            cur = cur.nexts[index];
        }
        cur.pass++;
        cur.end++;
    }

    /**
     * 从前缀树上删除一个字符串
     */
    public void delete(String word) {
        if (word == null) {
            return;
        }
        // 先查询看看该字符串是否存在于前缀树上
        if (search(word) > 0) {
            char[] chars = word.toCharArray();
            TrieNode cur = root;
            for (char aChar : chars) {
                int index = aChar - 'a';
                if (cur.nexts[index] == null) {
                    return;
                }
                cur.pass--;
                // help gc
                if (cur.pass <= 0) {
                    cur.nexts[index] = null;
                    return;
                }
                cur = cur.nexts[index];
            }
            cur.pass--;
            cur.end--;
        }
    }

    /**
     * 搜索一个字符串是否在前缀树上（是否被添加过）
     * 返回被添加的次数
     */
    public int search(String word) {
        if (word == null) {
            return 0;
        }
        char[] chars = word.toCharArray();
        TrieNode cur = root;
        for (char aChar : chars) {
            int index = aChar - 'a';
            TrieNode next = cur.nexts[index];
            if (next == null || next.pass <= 0) {
                return 0;
            }
            cur = cur.nexts[index];
        }
        return cur.end;
    }

    /**
     * 是否有以pre为前缀的字符串
     * 返回以pre为前缀的字符串数量
     */
    public int prefixNumber(String pre) {
        if (pre == null) {
            return 0;
        }
        TrieNode cur = root;
        char[] chars = pre.toCharArray();
        for (char aChar : chars) {
            int index = aChar - 'a';
            // 没有走向这个字符串的路
            if (cur.nexts[index] == null) {
                return 0;
            }
            cur = cur.nexts[index];
        }
        // 返回该节点被通过的次数
        return cur.pass;
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("wenpan");
        trie.insert("wenpanfeng");
        trie.insert("wenpanf");
        trie.insert("wenpanfe");
        trie.insert("wenpanfen");
        int wenpan = trie.search("wenpan");
        System.out.println("wenpan 是否存在：" + (wenpan > 0));
//        trie.delete("wenpan");
//        trie.delete("wenpanfeng");
//        trie.delete("wenpanf");
//        trie.delete("wenpanfen");
//        trie.delete("wenpanfe");

        String res = trie.search2("wen", 5);
        System.out.println("res = " + res);
        System.out.println("=================");
    }

}


